from tensorflow.keras.activations import relu
from tensorflow.keras.layers import (Activation, Add, BatchNormalization, Conv2D, DepthwiseConv2D)
from tensorflow.keras.models import Model


def make_divisible(v, divisor, min_value=None):
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v


def relu6(x):
    return relu(x, max_value=6)


class Mobile_block(Model):
    def __init__(self, expansion, stride, alpha, in_filters, filters, block_id, skip_connection, rate=1):
        super(Mobile_block, self).__init__()
        pointwise_conv_filters = int(filters * alpha)
        pointwise_filters = make_divisible(pointwise_conv_filters, 8)
        prefix = 'expanded_conv_{}_'.format(block_id)
        self.block_id = block_id
        self.skip_connection = skip_connection

        # ----------------------------------------------------#
        #       利用1x1卷积根据输入进来的通道数进行通道数上升
        # ----------------------------------------------------#
        if block_id:
            self.c1 = Conv2D(expansion * in_filters, kernel_size=1, padding='same', use_bias=False,
                             activation=None, name=prefix + 'expand')
            self.b1 = BatchNormalization(epsilon=1e-3, momentum=0.999, name=prefix + 'expand_BN')
            self.a1 = Activation(relu6, name=prefix + 'expand_relu')
        else:
            prefix = 'expanded_conv_'

        # ----------------------------------------------------#
        #             利用深度可分离卷积进行特征提取
        # ----------------------------------------------------#
        self.dc2 = DepthwiseConv2D(kernel_size=3, strides=stride, activation=None, use_bias=False,
                                   padding='same', dilation_rate=(rate, rate), name=prefix + 'depthwise')
        self.b2 = BatchNormalization(epsilon=1e-3, momentum=0.999, name=prefix + 'depthwise_BN')
        self.a2 = Activation(relu6, name=prefix + 'depthwise_relu')

        # ----------------------------------------------------#
        #              利用1x1的卷积进行通道数的下降
        # ----------------------------------------------------#
        self.c3 = Conv2D(pointwise_filters, kernel_size=1, padding='same', use_bias=False, activation=None,
                         name=prefix + 'project')
        self.b3 = BatchNormalization(epsilon=1e-3, momentum=0.999, name=prefix + 'project_BN')

        # ----------------------------------------------------#
        #                      添加残差边
        # ----------------------------------------------------#
        if skip_connection:
            self.add = Add(name=prefix + 'add')

    def call(self, inputs, training=None, mask=None):
        x = inputs
        if self.block_id:
            x = self.c1(x)
            x = self.b1(x)
            x = self.a1(x)
        x = self.dc2(x)
        x = self.b2(x)
        x = self.a2(x)
        x = self.c3(x)
        x = self.b3(x)
        if self.skip_connection:
            x = self.add([inputs, x])
        return x


class MobilenetV2(Model):
    def __init__(self, downsample_factor=8, alpha=1.0):
        super(MobilenetV2, self).__init__()
        if downsample_factor == 16:
            block4_dilation = 1
            block5_dilation = 2
            block4_stride = 2
        elif downsample_factor == 8:
            block4_dilation = 2
            block5_dilation = 4
            block4_stride = 1
        else:
            raise ValueError('Unsupported factor - `{}`, Use 8 or 16.'.format(downsample_factor))

        # Input(473,473,3)
        first_block_filters = make_divisible(32 * alpha, 8)
        # 473,473,3 -> 237,237,32
        self.mc1 = Conv2D(first_block_filters, kernel_size=3, strides=(2, 2), padding='same',
                          use_bias=False, name='Conv')
        self.mb1 = BatchNormalization(epsilon=1e-3, momentum=0.999, name='Conv_BN')
        self.ma1 = Activation(relu6, name='Conv_Relu6')

        # 237,237,32 -> 237,237,16
        self.mb2 = Mobile_block(in_filters=32, filters=16, alpha=alpha, stride=1, expansion=1,
                                block_id=0, skip_connection=False)

        # 237,237,16 -> 119,119,24
        self.mb3 = Mobile_block(in_filters=16, filters=24, alpha=alpha, stride=2, expansion=6,
                                block_id=1, skip_connection=False)
        self.mb4 = Mobile_block(in_filters=24, filters=24, alpha=alpha, stride=1, expansion=6,
                                block_id=2, skip_connection=True)

        # 119,119,24 -> 60,60.32
        self.mb5 = Mobile_block(in_filters=24, filters=32, alpha=alpha, stride=2, expansion=6,
                                block_id=3, skip_connection=False)
        self.mb6 = Mobile_block(in_filters=32, filters=32, alpha=alpha, stride=1, expansion=6,
                                block_id=4, skip_connection=True)
        self.mb7 = Mobile_block(in_filters=32, filters=32, alpha=alpha, stride=1, expansion=6,
                                block_id=5, skip_connection=True)

        # 60,60,32 -> 30,30.64
        self.mb8 = Mobile_block(in_filters=32, filters=64, alpha=alpha, stride=block4_stride,
                                expansion=6, block_id=6, skip_connection=False)
        self.mb9 = Mobile_block(in_filters=64, filters=64, alpha=alpha, stride=1, rate=block4_dilation,
                                expansion=6, block_id=7, skip_connection=True)
        self.mb10 = Mobile_block(in_filters=64, filters=64, alpha=alpha, stride=1, rate=block4_dilation,
                                 expansion=6, block_id=8, skip_connection=True)
        self.mb11 = Mobile_block(in_filters=64, filters=64, alpha=alpha, stride=1, rate=block4_dilation,
                                 expansion=6, block_id=9, skip_connection=True)

        # 30,30.64 -> 30,30.96
        self.mb12 = Mobile_block(in_filters=64, filters=96, alpha=alpha, stride=1, rate=block4_dilation,
                                 expansion=6, block_id=10, skip_connection=False)
        self.mb13 = Mobile_block(in_filters=96, filters=96, alpha=alpha, stride=1, rate=block4_dilation,
                                 expansion=6, block_id=11, skip_connection=True)
        self.mb14 = Mobile_block(in_filters=96, filters=96, alpha=alpha, stride=1, rate=block4_dilation,
                                 expansion=6, block_id=12, skip_connection=True)

        # 30,30.96 -> 30,30,160 -> 30,30,320
        self.mb15 = Mobile_block(in_filters=96, filters=160, alpha=alpha, stride=1, rate=block4_dilation,
                                 expansion=6, block_id=13, skip_connection=False)
        self.mb16 = Mobile_block(in_filters=160, filters=160, alpha=alpha, stride=1, rate=block5_dilation,
                                 expansion=6, block_id=14, skip_connection=True)
        self.mb17 = Mobile_block(in_filters=160, filters=160, alpha=alpha, stride=1, rate=block5_dilation,
                                 expansion=6, block_id=15, skip_connection=True)
        self.mb18 = Mobile_block(in_filters=160, filters=320, alpha=alpha, stride=1, rate=block5_dilation,
                                 expansion=6, block_id=16, skip_connection=False)

    def call(self, input):
        x = self.mc1(input)
        x = self.mb1(x)
        x = self.ma1(x)
        x = self.mb2(x)
        x = self.mb3(x)
        x = self.mb4(x)
        x = self.mb5(x)
        x = self.mb6(x)
        x = self.mb7(x)
        x = self.mb8(x)
        x = self.mb9(x)
        x = self.mb10(x)
        x = self.mb11(x)
        x = self.mb12(x)
        x = self.mb13(x)
        f4 = self.mb14(x)
        x = self.mb15(f4)
        x = self.mb16(x)
        x = self.mb17(x)
        f5 = self.mb18(x)
        return f4, f5


model = MobilenetV2()
model.build((None, 473, 473, 3))
model.summary()
